import time
from typing import List

import pandas as pd
import akshare as ak
import model.stock_index as stock_index


def get_index_list():
    index_list = []
    # 获取自选指数信息
    res = "./resource/指数列表.csv"
    df = pd.read_csv(res)
    for row in df.itertuples(index=False):
        code = getattr(row, '指数代码')
        dot_index = code.find('.')
        if dot_index != -1:
            code = code[0:dot_index]
        index_list.append(stock_index.StockIndex(code, getattr(row, '指数名称')))
    return index_list


def get_comm_index(a: List, b: List) -> int:
    rs = 0
    for e in a:
        for e2 in b:
           if e == e2:
                rs += 2
                break
    return rs


# 第一步 获取指数成分
idxMap = {}
idxList = get_index_list()
idxNames = []
for e in idxList:
    df = ak.index_stock_cons_weight_csindex(e.code)
    idxMap[e.code] = df['成分券代码'].tolist()
    idxNames.append(e.name)

data = {}
for e in idxList:
    percents = []
    for e2 in idxList:
        total = len(idxMap[e.code]) + len(idxMap[e2.code])
        comm = get_comm_index(idxMap[e.code], idxMap[e2.code])
        percents.append(round(comm * 100 / total))
    data[e.name] = percents

df = pd.DataFrame(index=idxNames, data=data)
print(df)
date = time.strftime("%Y%m%d", time.localtime())
df.to_excel(f'指数相似度-自选指数-{date}.xlsx')